#undef UNICODE

#include <Windows.h>
#include <stdio.h>
#include <stdlib.h>
#include <tchar.h>
#include <TlHelp32.h>
#include <winsock.h>

#include "common.h"
#include "serverSide.h"

#pragma comment(lib, "Ws2_32.lib")

int gSUid;
unsigned short gServerTcpPort;

SOCKET gUdpSd;
SOCKET gTcpSd;

#define SERVER_TCP_PORT_RANGE_BEGIN 50000
#define SERVER_TCP_PORT_RANGE_END	60000


void terminateServer()
{
	//use mutex

	if (gUdpSd != INVALID_SOCKET)
	{
		closesocket(gUdpSd);
	}

	if (gTcpSd != INVALID_SOCKET)
	{
		closesocket(gTcpSd);
	}

	WSACleanup();

	exit(0);
}





void serverHelloWorker()
{
	struct sockaddr_in client;

#pragma warning(suppress:4127) //this worker must live forever and serve TCP requests
	while(TRUE)
	{
		if (waitForClientHello(gUdpSd,gSUid,&client))
		{
			sendServerHello(gUdpSd,gServerTcpPort,gSUid,&client);
		}
	}
}


BOOL analyzeDirectoryContent(char* pDirPath,size_t* pTotalLength,unsigned short* pTotalCount)
{
	WIN32_FIND_DATA fileData;

	HANDLE fh = FindFirstFile(pDirPath,&fileData);

	size_t length;

    if (fh == INVALID_HANDLE_VALUE) 
    {
      printf ("[-] Invalid directory content request\n");
      return FALSE;
    }  
	do
	{
		length = strlen(fileData.cFileName);

		if ((length == 1 && fileData.cFileName[0] == '.') ||
			(length == 2 && strncmp(fileData.cFileName,"..",2)==0))
		{
			continue;
		}

		printf("%s (%d)\n",fileData.cFileName,strlen(fileData.cFileName));
		
		(*pTotalLength)+=strlen(fileData.cFileName);
		(*pTotalCount)++;
	}while(FindNextFile(fh, &fileData));
	FindClose(fh);
	return TRUE;
}

BOOL createFileListResponse(char* pBuffer,char* pDirPath,unsigned short numberOfElements)
{
	WIN32_FIND_DATA fileData;

	HANDLE fh = FindFirstFile(pDirPath,&fileData);

	size_t length;

    if (fh == INVALID_HANDLE_VALUE) 
    {
      printf ("[-] unable to create response\n");
      return FALSE;
    }  

	MEM_TO_VAL(pBuffer,unsigned short) = htons(PROTOCOL_PREFIX);
	MEM_TO_VAL(pBuffer+2,char) = DIR_CONTENT_RES_MSG_ID;
	MEM_TO_VAL(pBuffer+3,unsigned short) = htons(numberOfElements);

	pBuffer+=5;


	do
	{
		length = strlen(fileData.cFileName);

		if ((length == 1 && fileData.cFileName[0] == '.') ||
			(length == 2 && strncmp(fileData.cFileName,"..",2)==0))
		{
			continue;
		}

		printf("%s (%d)\n",fileData.cFileName,strlen(fileData.cFileName));
		
		MEM_TO_VAL(pBuffer,char) = fileData.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY; 
		MEM_TO_VAL(pBuffer+1,size_t) = htonl(fileData.nFileSizeLow);
		MEM_TO_VAL(pBuffer+5,unsigned short) = htons((unsigned short)length);
		memcpy(pBuffer+7,fileData.cFileName,length);

		pBuffer+=(7+length);

		
	}while(FindNextFile(fh, &fileData));
	FindClose(fh);
	return TRUE;

}

void handleFileListRequest(SOCKET sd,unsigned short pathLength)
{
	
	char* pDirPath = (char*)malloc((pathLength+3)*sizeof(char));
	char* pResponse = NULL;
	size_t totalLength = 0,responseSize=0;
	unsigned short elementsCount=0;

	do
	{
		if (pDirPath == NULL)
		{
			printf("[-] memory allocaion error\n");
			break;
		}

		memset(pDirPath,0x00,pathLength+3);
	
		if (recvTcpBuffer(sd,pDirPath,pathLength) != pathLength)
		{
			printf("[-] Error: can't read client request\n");
			break;
		}

		printf("[+] asking for directory %s content\n",pDirPath);

		if (pDirPath[pathLength-1]=='\\')
		{
			pDirPath[pathLength] = '*';
		}
		else
		{
			pDirPath[pathLength] = '\\';
			pDirPath[pathLength+1] = '*';
		}

		if (!analyzeDirectoryContent(pDirPath,&totalLength,&elementsCount))
		{
			break;
		}
		
		responseSize = (5+elementsCount*7 + totalLength)*sizeof(char);

		pResponse = (char*)malloc(responseSize);

		if (pResponse == NULL)
		{
			printf("[-] memory allocation error\n");
			return;
		}

		if (!createFileListResponse(pResponse,pDirPath,elementsCount))
		{
			break;
		}

		if (sendTcpBuffer(sd,pResponse,responseSize) != (int)responseSize)
		{
			printf("[-] unable to send response\n");
			break;
		}

#pragma warning(suppress:4127) //enable single exit point
	}while(FALSE);

	if (pDirPath != NULL)
	{
		free(pDirPath);
	}

	if (pResponse != NULL)
	{
		free(pResponse);
	}
}

void handleFileContentRequest(SOCKET sd,unsigned short fileNameLength)
{
	char* pFileName = (char*)malloc((fileNameLength+1)*sizeof(char));
	
	
	
	
	do
	{
		memset(pFileName,0x00,fileNameLength+1);
	
		if (recvTcpBuffer(sd,pFileName,fileNameLength) != fileNameLength)
		{
			printf("[-] Error: can't read client request\n");
			break;
		}

		printf("[+] asking for download %s\n",pFileName);
		
		sendFileContent(sd,pFileName);
		
		

#pragma warning(suppress:4127) //enable single exit point		
	}while(FALSE);

	if (pFileName != NULL);
	{
		free(pFileName);
	}

	

}

void handleTCPConnection(LPVOID param)
{
	SOCKET sd = (SOCKET)param;
	char clientMsgPrefix[5];

	//read msg header
	if (recv(sd,(char*)clientMsgPrefix,5,0) != 5)
	{
		closesocket(sd);
		return;
	}

	//decide that this message is
	if (IS_FILE_LIST_REQUEST(clientMsgPrefix))
	{
		handleFileListRequest(sd,ntohs(MEM_TO_VAL(clientMsgPrefix+3,unsigned short)));
	}
	else if (IS_FILE_CONTENT_REQUEST(clientMsgPrefix))
	{
		handleFileContentRequest(sd,ntohs(MEM_TO_VAL(clientMsgPrefix+3,unsigned short)));
	}

	closesocket(sd);

}

BOOL serverMainWorker()
{
	struct sockaddr_in client;
	int client_length = sizeof(struct sockaddr_in);
	
	SOCKET newConnectionSd;

#pragma warning(suppress:4127) //this thread must live forever and serve TCP requests
	while(TRUE)
	{
		newConnectionSd = accept(gTcpSd,(sockaddr*)&client,&client_length);
		
		if (newConnectionSd == SOCKET_ERROR)
		{
			continue;
		}
		
		if (CreateThread(0,0,(LPTHREAD_START_ROUTINE)&handleTCPConnection,(LPVOID)newConnectionSd,0,NULL) == NULL)
		{
			printf("Error: unable to create TCP handler thread\n");
			break;
		}

	}

	return FALSE;

	
}


int _tmain(int argc, _TCHAR* argv[])
{
	WSADATA w;					

	if (argc != 2)
	{
		_tprintf(_T("Usage: %s <SID>\n"),argv[0]);
		exit(-1);
	}

	do
	{	
		gSUid =  _wtoi(argv[1]);
	

		if (gSUid == 0)
		{
			printf("Error: invalid <SID>\n");
			exit(-1);
		}

		if (WSAStartup(0x0101, &w) != 0)
		{
			printf("Error: Could not open Windows connection.\n");
			break;
		}
	
		gUdpSd = createUDPSocket(SERVER_UDP_PORT,FALSE);

		if (gUdpSd == SOCKET_ERROR)
		{
			printf("Error: unable to create UDP socket\n");
			break;
		}

		
		gServerTcpPort = SERVER_TCP_PORT_RANGE_BEGIN-1;

		do
		{
			++gServerTcpPort;
			gTcpSd = createServerTCPSocket(gServerTcpPort);
		}while(gTcpSd == SOCKET_ERROR && gServerTcpPort < SERVER_TCP_PORT_RANGE_END);


		

		if (gTcpSd == SOCKET_ERROR)
		{
			printf("Error: unable to create TCP socket\n");
			break;
		}

		if (CreateThread(NULL,0,(LPTHREAD_START_ROUTINE)&serverHelloWorker,NULL,0,NULL)==NULL)
		{
			break;
		}


		if (!serverMainWorker())
		{
			break;
		}

#pragma warning(suppress:4127) //enable single exit point	
	}while(FALSE);

	terminateServer();

	return -1;
}